import os
import pandas as pd
import numpy as np
from matplotlib.figure import Figure
from matplotlib import gridspec
from sqlalchemy import create_engine
from WindPy import w


def download_from_wind(name_dic):
    """
    根据wind_id下载数据，并生成DataFrame，时间默认从2016年开始
    :param name_dic: 要包括 'id' 这个key，目前只能下载一个
    :return: DataFrame
    """
    data_id = [name_dic['id']]
    if not w.isconnected():
        w.start(showmenu=False)
    resp = w.edb(data_id, beginTime='2016-01-01')
    da = pd.DataFrame(resp.Data, columns=resp.Times, index=resp.Codes).T
    da.index = pd.to_datetime(da.index)
    return da


def plot_and_seasonal_plot(name_dic, fig):
    name_dir = f"data/{name_dic['name']}"
    if not os.path.exists(name_dir):
        os.makedirs(name_dir)
    else:
        print("exists", name_dir)

    save_path = f"{name_dir}/{name_dic['save_name']}.png"

    if name_dic['get'] == 'wind':
        data = download_from_wind(name_dic)
    elif name_dic['get'] == 'func':
        func = eval(name_dic['func'])
        data = func(name_dic)

    # fig = Figure(figsize=(20, 12), dpi=100)
    spec = gridspec.GridSpec(ncols=4, nrows=4, figure=fig)
    ax1 = fig.add_subplot(spec[:2, :2])
    ax2 = fig.add_subplot(spec[:2, 2:])

    se = data.iloc[:, 0]
    ax1.plot(se, label=f"{name_dic['name'] + name_dic['attr']}")

    dt = pd.DataFrame(se)
    dt['t'] = dt.index
    dt['year'] = dt['t'].dt.year
    dt['day_of_year'] = dt['t'].dt.dayofyear
    for year, year_df in dt.groupby(by='year'):
        year_df.index = year_df['day_of_year']
        ax2.plot(year_df.iloc[:, 0], label=year)
    ax2.legend()

    for ax in [ax1, ax2]:
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)

    # fig.savefig(save_path, bbox_inches="tight")
    print('save >>', save_path)

    return fig


def get_future_spot_basis(name_dic):
    print(name_dic)
    db_future_main = 'mysql+pymysql://hf_user:290202@192.168.3.179:3306/future?charset=utf8&use_unicode=1'
    engine_future_main = create_engine(db_future_main)
    spot_se = download_from_wind(name_dic).iloc[:, 0]
    future_data = pd.read_sql(f"select * from {name_dic['dbname']}", engine_future_main)
    future_data.index = pd.to_datetime([str(s) for s in future_data.trade_date])
    future_se = future_data.close
    new_se_idx = [s for s in future_se.index if s in spot_se.index]
    future_se = future_se[new_se_idx]  # 设定同一索引
    spot_se = spot_se[new_se_idx]
    new_se = spot_se - future_se  # 现货 - 期货
    new_se.name = name_dic['name'] + name_dic['save_name']
    return pd.DataFrame(new_se)


def get_future_future_basis(name_dic):
    db_future_main = 'mysql+pymysql://hf_user:290202@192.168.3.179:3306/future_short?charset=utf8&use_unicode=1'
    engine_future_ = create_engine(db_future_main)
    tbs = pd.read_sql("show tables", engine_future_).iloc[:,0].to_list()
    tk, tp = name_dic['dbname'].split('_')
    first_data = get_month_future(tk, tp, engine_future_, tbs, '01')
    second_data = get_month_future(tk, tp, engine_future_, tbs, '05')
    new_se_idx = [s for s in first_data.index if s in second_data.index]
    first_data = first_data[new_se_idx]
    second_data = second_data[new_se_idx]
    basis = first_data - second_data
    basis.name = f"{name_dic['name']}01-05"
    return pd.DataFrame(basis)


def get_month_future(tk, tp, engine, tbs, month_str):
    start = True
    for year in [17, 18, 19, 20, 21, 22]:
        db_first = f"{tk.lower()}{year}{month_str}_{tp.lower()}"
        sql_first = f"select * from {db_first}"
        print(sql_first)
        assert db_first in tbs
        temp_data_first = pd.read_sql(sql_first, engine)
        temp_data_first.index = [pd.to_datetime(str(s)) for s in temp_data_first.trade_date]
        temp_data_first.sort_index(inplace=True, ascending=True)
        if start:
            result_data_first = temp_data_first.close
            start = False
        else:
            result_data_first = pd.concat([result_data_first, temp_data_first.close])
    return result_data_first


if __name__ == "__main__":
    # 期限基差绘图参数
    name_dic_fs_bs = {"name": "动力煤", "id_name": "主力合约", "attr": "期货", "id": "S5120089", "save_name": '基差',
                      'get': 'func', 'dbname': 'ZC_ZCE', 'func': 'get_future_spot_basis'}
    # 期货价差绘图参数
    name_dic_ff_bs = {"name": "动力煤", "id_name": "主力合约", "attr": "期货", "id": "-", "save_name": '价差',
                      'get': 'func', 'dbname': 'ZC_ZCE', 'func': 'get_future_future_basis'}
    # 现货价格绘图参数
    name_dic_1 = {"name": "动力煤", "id_name": "市场价:动力煤(Q5500,山西产):秦皇岛", "attr": "现货",
                  "id": "S5120089", "save_name": '现货', 'get': 'wind'}
    # 库存绘图参数
    name_dic_2 = {"name": "动力煤", "id_name": "CCTD主流港口合计煤炭库存", "attr": "库存",
                  "id": "S5134687", "save_name": '库存', 'get': 'wind'}

    names = [name_dic_fs_bs, name_dic_ff_bs, name_dic_1, name_dic_2]
    for name in names:
        plot_and_seasonal_plot(name)

